Hi @giulio98,

Thanks for your contribution. I started looking at the PR. Two initial remarks:

1. Your PR should not update pipeline.py, in particular because it silently concatenate context and question tokens. We'd rather prefer to modify evaluate.py as you initially suggested (adding a [SEP] special token)
2. Your PR should not modify qfilter_press.py, duo_attention_press.py and longbech/calculate_metrics.py

I started to look more into finch_press.py. Would it be possible to re-use components from other presses ? For instance you could delete compute_finch_attention and compute_normalization_factors and simply use

    def score(self, module, hidden_states, keys, values, attentions, kwargs):

        bsz, num_key_value_heads, q_len, _ = keys.shape
        num_key_value_groups = module.config.num_attention_heads // num_key_value_heads

        if attentions is not None:
            attn_weights = attentions[..., -self.condition_len :, : -self.condition_len]
        else:
            attn_weights = SnapKVPress.compute_window_attention(module, hidden_states, keys, self.condition_len, kwargs["position_embeddings"])

        if self.normalize_scores:
            non_zero_counts = torch.arange(q_len - self.condition_len, q_len)[None, None, :, None]
            non_zero_counts = non_zero_counts.to(attn_weights.device)
            attn_weights = attn_weights * non_zero_counts

        # Average per group
        scores = attn_weights.mean(dim=-2)
        scores = scores.view(bsz, num_key_value_heads, num_key_value_groups, q_len - self.condition_len)
        scores = scores.mean(dim=2)
        
        # Add back the observation window. Use max score to make sure the window is not pruned.
        scores = F.pad(scores, (0, self.condition_len), value=scores.max().item())
        return scores

Comparison with the initial score method and the score method I propose:

Note that in the code above I replaced sum by mean to avoid too large floats when using bfloat16. Also I replaced compute_normalization_factors by a simple torch.arange which is equivalent. Similarly, I will look if it's possible to use ChunkPress and KeyRerotationPress.

Hello @SimJeg
I deleted my modifications for longbench, duo_filter_press and qfilter_press it was modified automatically when I run black.
Regarding the score I applied your suggestion, however I keep the call to compute_finch_attention, because when split_size is different than 1, the diagonal for attention mask in line 70 has to be computed like this:

attention_mask = torch.triu(attention_mask, diagonal=key_states.shape[-2] - condition_len + 1)

instead of

attention_mask = torch.triu(attention_mask, diagonal=q_len - window_size + 1)

because it needs into account for the past cache of previous iteration, similarly I had to overwrite forward for the following line:

        q_len = hidden_states.shape[1]

        # Don't compress after pre-filling
        if kwargs["cache_position"][-1] > q_len + cache.get_seq_length():
            return output

line 190-194.

For same reason I couldn't use KeyRerotationPress because on the way the nkept is calculated that has to take into account the previous cache length.

Finally, ChunkPress doesn't implement the chunked forward correctly, we wrote for this pourpose the chunked_forward.

Regarding the special [SEP] token, I thought it will make a bit hacky for the user to use FinchPress, but I'm open to discussion for a better user friendly usage.

Giulio

In SnapKVPress we have

bsz, num_key_value_heads, q_len, _ = keys.shape

so in fact q_len is already keys.shape[-2] and you can remove compute_finch_attention.

In SnapKVPress we have

bsz, num_key_value_heads, q_len, _ = keys.shape

so in fact q_len is already keys.shape[-2] and you can remove compute_finch_attention.

You are right I replaced it with SnapKVPress compute window attention

@giulio98 do you think it would make a big difference (in terms of accuracy) to move from :

1. Run forward pass for x=cat([chunk, question])
2. Compute scores, compress and rerotate KV cache for chunk
3. Go to next chunk

to

1. Run forward pass for the whole inputs
2. Compute scores by chunk (as in ChunkPress)
3. Compress and rerotate (as in KeyRerotationPress)

I know it's different but the code would be much easier to read (which is my main concern so far), and maybe have similar performances

@giulio98 do you think it would make a big difference (in terms of accuracy) to move from :

1. Run forward pass for x=cat([chunk, question])
2. Compute scores, compress and rerotate KV cache for chunk
3. Go to next chunk

to

1. Run forward pass for the whole inputs
2. Compute scores by chunk (as in ChunkPress)
3. Compress and rerotate (as in KeyRerotationPress)

I know it's different but the code would be much easier to read (which is my main concern so far), and maybe have similar performances

My main concern about this implementation is that the chunked forward was designed in Finch explicttly to handle inputs that exceeed the context window size of LLMs (e.g 128k tokens) using the chunked forward input is split into managable chunks whose length is less than the context window size and that can be processed correctly by the model.

@giulio98 do you think it would make a big difference (in terms of accuracy) to move from :

1. Run forward pass for x=cat([chunk, question])
2. Compute scores, compress and rerotate KV cache for chunk
3. Go to next chunk

to

1. Run forward pass for the whole inputs
2. Compute scores by chunk (as in ChunkPress)
3. Compress and rerotate (as in KeyRerotationPress)

I know it's different but the code would be much easier to read (which is my main concern so far), and maybe have similar performances

My main concern about this implementation is that the chunked forward was designed in Finch explicttly to handle inputs that exceeed the context window size of LLMs (e.g 128k tokens) using the chunked forward input is split into managable chunks whose length is less than the context window size and that can be processed correctly by the model.

On the other hand I know that doing as you proposed will make easier to apply for example AdaKV on Finch with possible performance gains

Indeed with a compression ratio of 50%, Finch could handle inputs up to 256k even if the LLM max size is 128k (and more generally max_len / compression_ratio). However, even on RULER with 4k (which is far below 128k) we see degradations at 50%, so it might not be relevant to focus on the 128k regime, WDYT ?

I suspect that RULER is a dataset very information dense, so it is very hard to apply compression beyond a certain limit, instead if we imagine to compress external knowledge from let's say wikipedia pages I think will make a lot of sense KV Compression.

Draft for a proposal of a simplified FinchPress:

# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0


import math
from dataclasses import dataclass

import torch
from torch.nn import functional as F

from kvpress.presses.base_press import BasePress
from kvpress.presses.snapkv_press import SnapKVPress
from transformers.models.llama.modeling_llama import rotate_half


@dataclass
class FinchPress(BasePress):

    compression_ratio: float = 0.0
    split_size: int = 1
    normalize_scores: bool = True
    condition_len: int = None

    def score(self, module, hidden_states, keys, values, attentions, kwargs):
        """
        Similar to SnapKVPress except it adds a normalization step before averaging on the context window.
        """

        bsz, num_key_value_heads, q_len, _ = keys.shape
        num_key_value_groups = module.config.num_attention_heads // num_key_value_heads

        if attentions is not None:
            attn_weights = attentions[..., -self.condition_len :, : -self.condition_len]
        else:
            attn_weights = SnapKVPress.compute_window_attention(
                module, hidden_states, keys, self.condition_len, kwargs["position_embeddings"]
            )

        if self.normalize_scores:
            non_zero_counts = torch.arange(q_len - self.condition_len, q_len)[None, None, :, None]
            non_zero_counts = non_zero_counts.to(attn_weights.device)
            attn_weights = attn_weights * non_zero_counts

        # Average per group
        scores = attn_weights.mean(dim=-2)
        scores = scores.view(bsz, num_key_value_heads, num_key_value_groups, q_len - self.condition_len)
        scores = scores.mean(dim=2)

        # Add back the observation window. Use max score to make sure the window is not pruned.
        scores = F.pad(scores, (0, self.condition_len), value=scores.max().item())
        return scores

    def compress(self, module, hidden_states, keys, values, attentions, kwargs):
        """
        Scores are computed by chunks, keys and values are then compressed and re-rotated.
        """

        q_len = hidden_states.shape[1]

        if self.compression_ratio == 0:
            return keys, values
        assert (self.condition_len is not None) and (self.condition_len < q_len)

        # Compute scores
        scores = self.score(module, hidden_states, keys, values, attentions, kwargs)

        # Compute indices by chunks
        indices = []
        chunk_size = math.ceil(q_len / self.split_size)
        for i, chunk_scores in enumerate(torch.split(scores, chunk_size, dim=2)):
            n_kept = max(1, int(chunk_scores.shape[2] * (1 - self.compression_ratio)))
            chunk_indices = i * chunk_size + chunk_scores.topk(n_kept, dim=-1).indices
            indices.append(chunk_indices)

        indices = torch.cat(indices, dim=-1)
        indices = indices.unsqueeze(-1).expand(-1, -1, -1, module.head_dim)

        # Rerotate keys and values
        cos, sin = kwargs["position_embeddings"]
        keys = (keys * cos.unsqueeze(1)) + (rotate_half(keys) * (-sin.unsqueeze(1)))
        keys = keys.gather(2, indices).contiguous()
        cos, sin = cos[:, : indices.shape[2]], sin[:, : indices.shape[2]]
        keys = (keys * cos.unsqueeze(1)) + (rotate_half(keys) * sin.unsqueeze(1))

        values = values.gather(2, indices).contiguous()

        return keys, values

I will run it for 50% compression to compare with what you reported

@SimJeg
Also, I noticed that for small compression rate (e.g 0.1, 0.25) FinchPress has degraded performance with respect to SnapKVPress which is a bit weird, maybe the inclusion of the answer_prefix in the condition degrate performance?
Would you happen to have compute available to run a quick experiment to test this hypothesis?

Draft for a proposal of a simplified FinchPress:

# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0


import math
from dataclasses import dataclass

import torch
from torch.nn import functional as F

from kvpress.presses.base_press import BasePress
from kvpress.presses.snapkv_press import SnapKVPress
from transformers.models.llama.modeling_llama import rotate_half


@dataclass
class FinchPress(BasePress):

    compression_ratio: float = 0.0
    split_size: int = 1
    normalize_scores: bool = True
    condition_len: int = None

    def score(self, module, hidden_states, keys, values, attentions, kwargs):
        """
        Similar to SnapKVPress except it adds a normalization step before averaging on the context window.
        """

        bsz, num_key_value_heads, q_len, _ = keys.shape
        num_key_value_groups = module.config.num_attention_heads // num_key_value_heads

        if attentions is not None:
            attn_weights = attentions[..., -self.condition_len :, : -self.condition_len]
        else:
            attn_weights = SnapKVPress.compute_window_attention(
                module, hidden_states, keys, self.condition_len, kwargs["position_embeddings"]
            )

        if self.normalize_scores:
            non_zero_counts = torch.arange(q_len - self.condition_len, q_len)[None, None, :, None]
            non_zero_counts = non_zero_counts.to(attn_weights.device)
            attn_weights = attn_weights * non_zero_counts

        # Average per group
        scores = attn_weights.mean(dim=-2)
        scores = scores.view(bsz, num_key_value_heads, num_key_value_groups, q_len - self.condition_len)
        scores = scores.mean(dim=2)

        # Add back the observation window. Use max score to make sure the window is not pruned.
        scores = F.pad(scores, (0, self.condition_len), value=scores.max().item())
        return scores

    def compress(self, module, hidden_states, keys, values, attentions, kwargs):
        """
        Scores are computed by chunks, keys and values are then compressed and re-rotated.
        """

        q_len = hidden_states.shape[1]

        if self.compression_ratio == 0:
            return keys, values
        assert (self.condition_len is not None) and (self.condition_len < q_len)

        # Compute scores
        scores = self.score(module, hidden_states, keys, values, attentions, kwargs)

        # Compute indices by chunks
        indices = []
        chunk_size = math.ceil(q_len / self.split_size)
        for i, chunk_scores in enumerate(torch.split(scores, chunk_size, dim=2)):
            n_kept = max(1, int(chunk_scores.shape[2] * (1 - self.compression_ratio)))
            chunk_indices = i * chunk_size + chunk_scores.topk(n_kept, dim=-1).indices
            indices.append(chunk_indices)

        indices = torch.cat(indices, dim=-1)
        indices = indices.unsqueeze(-1).expand(-1, -1, -1, module.head_dim)

        # Rerotate keys and values
        cos, sin = kwargs["position_embeddings"]
        keys = (keys * cos.unsqueeze(1)) + (rotate_half(keys) * (-sin.unsqueeze(1)))
        keys = keys.gather(2, indices).contiguous()
        cos, sin = cos[:, : indices.shape[2]], sin[:, : indices.shape[2]]
        keys = (keys * cos.unsqueeze(1)) + (rotate_half(keys) * sin.unsqueeze(1))

        values = values.gather(2, indices).contiguous()

        return keys, values

I will run it for 50% compression to compare with what you reported

@SimJeg I just noticed that you may have to order the indices as in line 114 in finch_press.py, as we saw this will enhance performance.

@giulio98 I created a new branch here: https://github.com/NVIDIA/kvpress/tree/simon/finch

What it contains:

• improved version of my draft above
• new way to compute condition_len (renamed window_size to align with SnapKV) which avoid important refacto of the pipeline with the silent concatenation of context_ids and question_ids. This is inspired by what you first suggested in the issue (FinchPress Scorer #59). I used <bos_token> instead of a new [SEP] token as it's possible to access it directly from the model (so no need to add extra param to the press).
• moved from split_size to chunk_length to align with the paper

I will share results with 50% compression. Could you provide the detailed performances for each subtask you obtained ?

update:

cwe= 94.29
fwe= 91.03
niah_multikey_1= 100.0
niah_multikey_2= 94.74
niah_multikey_3= 88.89
niah_multiquery= 97.5
niah_multivalue= 96.67
niah_single_1= 100.0
niah_single_2= 100.0
niah_single_3= 92.59
qa_1= 79.17
qa_2= 40.74
vt= 91.76

Average: 89.80

So it's slightly lower than your version. Once slight difference however is that I compress the question too but might impact only ~0.5% of compression ratio.

@giulio98 I created a new branch here: https://github.com/NVIDIA/kvpress/tree/simon/finch

What it contains:

• improved version of my draft above

• new way to compute condition_len (renamed window_size to align with SnapKV) which avoid important refacto of the pipeline with the silent concatenation of context_ids and question_ids. This is inspired by what you first suggested in the issue (FinchPress Scorer #59). I used <bos_token> instead of a new [SEP] token as it's possible to access it directly from the model (so no need to add extra param to the press).

• moved from split_size to chunk_length to align with the paper

I will share results with 50% compression. Could you provide the detailed performances for each subtask you obtained ?

update:

cwe= 94.29

fwe= 91.03

niah_multikey_1= 100.0

niah_multikey_2= 94.74

niah_multikey_3= 88.89

niah_multiquery= 97.5

niah_multivalue= 96.67

niah_single_1= 100.0

niah_single_2= 100.0

niah_single_3= 92.59

qa_1= 79.17

qa_2= 40.74

vt= 91.76



Average: 89.80

So it's slightly lower than your version. Once slight difference however is that I compress the question too but might impact only ~0.5% of compression ratio.

Hello,

I have to rerun the experiment, meantime first thing I noticed is missing is the sorting of the indices just after the topk, in fact top k return indices according to higher score by default, however we may need to sort them to their natural order because otherwise they can assume different meaning (this can be one thing that can enhance performance also in other presses).

Great catch ! It also impacts KeyRerotationPress so I will correct it, but it won't impact other presses as the order of keys and values does not matter.

I re-ran your implementation and get 91.6. Also I made a mistake in the numbers I reported above (I reported results for 5% of the data). With the error, I get 90.8.

I will correct the error and report results for 4 options: with / without - normalization / rerotation.

Great catch ! It also impacts KeyRerotationPress so I will correct it, but it won't impact other presses as the order of keys and values does not matter.

I re-ran your implementation and get 91.6. Also I made a mistake in the numbers I reported above (I reported results for 5% of the data). With the error, I get 90.8.

I will correct the error and report results for 4 options: with / without - normalization / rerotation.

Ah yes! If no rerotation is applied it is permutation invariant.

Updated results. Fixing the bug led to slighly worse performance (especially for niah_multikey_3).
Also again, the original implementation slightly over-estimates scores because compression ratio is not exactly the correct one.
Results w/o re-rotation are not reported because I forgot it would imply to update pipeline.py too.

Will look again to what might be the difference.

Ok I found the difference, in your implementation the condition_len tokens include the question but also the end of the chat template and the answer prefix. My version is aligned with what we did with other presses (only include the question). So I would favor the new PR I opened.